#### 
import topology_engineer, traffic_engineer, network_controller, path_selector, traffic_history_reader, traffic_pb2, cluster_finder
from aurora_network import *
import matplotlib.pyplot as plt
import numpy as np
import random


## generates a random traffic matrix
def generate_random_point(nblocks, norm=1.):
	norm_leftover = float(norm)
	pairs_vector = []
	tm = np.zeros((nblocks, nblocks))
	for i in range(nblocks - 1):
		for j in range(i + 1, nblocks, 1):
			pairs_vector.append((i,j))
			pairs_vector.append((j,i))
	random.shuffle(pairs_vector)
	for (i,j) in pairs_vector:
		tm[i][j] = int(random.random() * norm_leftover)
		norm_leftover -= tm[i][j]
	i = pairs_vector[0][0]
	j = pairs_vector[0][1]
	tm[i][j] += int(norm_leftover)
	return tm

def transform_traffic_matrix_to_traffic_snapshot(aurora_network, traffic_matrix):
	snapshot = traffic_pb2.TrafficSnapshot()
	nblocks = aurora_network.get_num_blocks()
	for i in range(nblocks):
		src_name = aurora_network.id_to_names_dict[i]
		for j in range(nblocks):
			if i != j and traffic_matrix[i][j] > 0:
				dst_name = aurora_network.id_to_names_dict[j]
				traffic_record = snapshot.records.add()
				traffic_record.source = src_name
				traffic_record.destination = dst_name
				traffic_record.max_load_gbps = int(traffic_matrix[i][j])
	return snapshot

def write_to_file(filename, x, ave_toe_mlu, max_toe_mlu, robust_toe_mlu):
	assert(len(x) == len(ave_toe_mlu))
	assert(len(x) == len(max_toe_mlu))
	assert(len(x) == len(robust_toe_mlu))
	str_builder = "x ave_toe max_toe robust_toe \n"
	for i in range(len(x)):
		str_builder += "{} {} {} {} \n".format(x[i], ave_toe_mlu[i], max_toe_mlu[i], robust_toe_mlu[i])
	with open(filename, "w+") as f:
		f.write(str_builder)
	return

def main(nblocks, block_params, number_of_random_points, numK):
	# Step 2 : Contruct the physical network itself
	block_names_list = ["ju1_block{}".format(x) for x in range(nblocks)]
	aurora_network = AuroraNetwork("artificial", block_params, block_names_list)

	generated_traffic_matrices = []
	generated_traffic_snapshots = []
	for _ in range(number_of_random_points):
		generated_tm = generate_random_point(nblocks, norm=1E4)
		generated_traffic_matrices.append(generated_tm)
		generated_traffic_snapshots.append(transform_traffic_matrix_to_traffic_snapshot(aurora_network, generated_tm))

	all_interblock_paths = path_selector.PathSelector(aurora_network, use_multihop=True).get_all_paths()
	print all_interblock_paths
	print("aurora's nblocks is : {}".format(aurora_network.get_num_blocks()))
	## ToEs 
	ave_toe = topology_engineer.HistoricalAveTrafficTopologyEngineer(aurora_network, 5,
																	5, all_interblock_paths, 
																	[])
	max_toe = topology_engineer.HistoricalMaxTrafficTopologyEngineer(aurora_network, 5,
																	5, all_interblock_paths, 
																	[])
	robust_toe = topology_engineer.RobustMultiTrafficTopologyEngineer(aurora_network, 5,
																	5, all_interblock_paths, 
																	[], numK, ignore_flows=[])
	ave_topology = aurora_network.round_fractional_topology_giant_switch(ave_toe._topology_engineer(generated_traffic_snapshots, all_interblock_paths), [])
	max_topology = aurora_network.round_fractional_topology_giant_switch(max_toe._topology_engineer(generated_traffic_snapshots, all_interblock_paths), [])
	robust_topology = aurora_network.round_fractional_topology_giant_switch(robust_toe._topology_engineer(generated_traffic_snapshots, all_interblock_paths, spare_traffic=0.1, fraction_threshold=0.002), [])
	## TEs
	ave_te = traffic_engineer.HistoricalAveTrafficEngineer(aurora_network, all_interblock_paths, 5, 5)
	max_te = traffic_engineer.HistoricalMaxTrafficEngineer(aurora_network, all_interblock_paths, 5, 5)
	robust_te = traffic_engineer.RobustMultiClusterScaleUpWeightedTrafficEngineer(aurora_network, all_interblock_paths, 5, 5, numK, ignore_flows=[])

	ave_routing_weights = ave_te.compute_path_weights(ave_topology, generated_traffic_matrices)
	max_routing_weights = max_te.compute_path_weights(ave_topology, generated_traffic_matrices)

	training_traffic_points = [np.array(aurora_network.get_traffic_vector_from_traffic_snapshot(x, normalize_to=1.)) for x in generated_traffic_snapshots] 
	clustering_algo = cluster_finder.ClusteringAlgorithm(numK, nfeatures=-1)
	traffic_cluster_points = clustering_algo.compute_cluster_centroid(np.array(training_traffic_points))
	centroid_traffic_matrices = [aurora_network.normalize_matrix(aurora_network.transform_traffic_point_to_traffic_matrix(x), norm=1.) for x in traffic_cluster_points]
	robust_routing_weights = robust_te.compute_path_weights(robust_topology, centroid_traffic_matrices, spare_traffic=0.1, fraction_threshold=0.001)
	
	## form a dummy controller for evaluation
	dummy_netcontroller = network_controller.NetworkController(aurora_network, None, None, number_of_random_points)
	ave_mlu = []
	max_mlu = []
	robust_mlu = []
	offset = 0
	for traffic_matrix in generated_traffic_matrices:
		max_perf = dummy_netcontroller._evaluate_snapshot_performance(offset, traffic_matrix, max_topology, max_routing_weights)
		ave_perf = dummy_netcontroller._evaluate_snapshot_performance(offset, traffic_matrix, ave_topology, ave_routing_weights)
		robust_perf = dummy_netcontroller._evaluate_snapshot_performance(offset, traffic_matrix, robust_topology, robust_routing_weights)
		ave_mlu.append(ave_perf[0])
		max_mlu.append(max_perf[0])
		robust_mlu.append(robust_perf[0])
		offset += 1
	fig = plt.figure()
	plt.plot(sorted(robust_mlu))
	plt.plot(sorted(ave_mlu))
	plt.plot(sorted(max_mlu))
	plt.legend(["Robust ToE", "Ave ToE", "Max ToE"])
	plt.ylabel("MLU")
	fig = plt.figure()
	plt.plot(robust_mlu)
	plt.plot(ave_mlu)
	plt.plot(max_mlu)
	plt.legend(["Robust ToE", "Ave ToE", "Max ToE"])
	plt.ylabel("MLU")
	plt.show()


	#write_to_file("/Users/minyee/static_topology_engineer/plot_data/multi_traffic_toe_motivation.txt", range(len(generated_traffic_matrices)), ave_mlu, max_mlu, robust_mlu)
	#write_to_file("/Users/minyee/static_topology_engineer/plot_data/multi_traffic_toe_motivation_sorted.txt", range(len(generated_traffic_matrices)), sorted(ave_mlu), sorted(max_mlu), sorted(robust_mlu))
	return







if __name__ == "__main__":
	print("\n\n")
	print("############################################")
	print("#  Begins")
	print("############################################")
	numK = 4
	number_of_random_points = 30
	nblocks = 8
	superblock_num_links = 256
	border_router_num_links = 256
	superblock_link_capacity = 100
	border_router_link_capacity = 100

	block_params = {BlockType.SUPERBLOCK : {}, BlockType.BORDER_ROUTER : {}}
	block_params[BlockType.SUPERBLOCK]["link capacity"] = float(superblock_link_capacity)
	block_params[BlockType.SUPERBLOCK]["num links"] = float(superblock_num_links)
	block_params[BlockType.BORDER_ROUTER]["link capacity"] = float(border_router_link_capacity) # in gbps
	block_params[BlockType.BORDER_ROUTER]["num links"] = float(border_router_num_links)

	main(nblocks, block_params, number_of_random_points, numK)
	print("############################################")
	print("# Completes")
	print("############################################")
	print("\n\n")